import tensorflow as tf


class BAP(tf.keras.layers.Layer):
    def __init__(self, feature_shape, att_nums, *args, **kargs):
        super(BAP, self).__init__(*args, **kargs)

        self.feature_size = feature_shape[0:2]
        self.concat = tf.keras.layers.Concatenate(axis=1)
        self.att_size = feature_shape[2] * att_nums

    def call(self, features, attentions):
        EPSILON = 1e-5
        UH, UW = self.feature_size
        B, H, W, C = features.get_shape().as_list()
        _, AH, AW, M = attentions.get_shape().as_list()
        assert UH == H and UW == W and UH == AH and UW == AW

        temp_feat = tf.einsum('ijkm,ijkn->imn', attentions, features) / float(H * W)
        feature_matrix = tf.reshape(temp_feat, (-1, self.att_size))

        feature_matrix = tf.sign(feature_matrix) * tf.sqrt(tf.abs(feature_matrix) + EPSILON)

        feature_matrix = tf.math.l2_normalize(feature_matrix)
        return feature_matrix


if __name__ == '__main__':
    model = BAP((7, 7, 960), 4)

    x = tf.random.normal((2, 7, 7, 960))
    att = tf.random.normal((2, 7, 7, 4))

    y = model(x, att)
    print(y)
